#-*- coding: utf-8 -*-

class StockModel: # {{{ 
    """
    Model danych giełdowych zawierający wszystkie dane o spółce 
    potrzebne strategii inwestycyjnej do analizy.

    Ponieważ wszystko poza typami wbudowanymi jest przekazywane przez referencję, obiekt 
    klasy StockData jest przekazywany każdej strategii jako jedyny argument,
    dzięki czemu unikamy kopiowania dużych danych
    """

    # Nie zmieniać samemu 
    firstDayIndex   =  0
    currentDayIndex = -1
    lastDayIndex    =  0

    # Wartość wskaźników dla bieżącego dnia. Gdy nie ma delty (tj. interesuje nas bieżący dzień)
    # to lepiej stosować te pola zamiast akcesorów (get*()), bo dzięki temu przy zapytaniu
    # jest jeden lookup w słowniku mniej

    openPrice   = None
    highPrice   = None
    lowPrice    = None
    closePrice  = None
    volume      = None
    OBV         = 0

    def reset(self, firstDay = None, lastDayIndex = None):
        if firstDay == None:
            firstDay =  self.firstDayIndex - 1

        self.currentDayIndex = self.firstDayIndex - 1
        self.advanceDay()

    def getStochasticOscillator(self, days):
        if self._Lowest[days][self.currentDayIndex] == -1:
              return -1
        liczba = self._Lowest[days][self.currentDayIndex]
        licznik = self.closePrice - self._Lowest[days][self.currentDayIndex]
        mianownik = (self._Highest[days][self.currentDayIndex] - self._Lowest[days][self.currentDayIndex])
        if mianownik == 0:
              return -1

        return licznik/mianownik

    def getSharePrice(self, *args):
        return       self.getClosePrice(*args)

    def getOpenPrice(self, *args):
        return float(self._getColumn(2, *args))

    def getHighPrice(self, *args):
        return float(self._getColumn(3, *args))

    def getLowPrice(self, *args):
        return float(self._getColumn(4, *args))

    def getClosePrice(self, *args):
        return float(self._getColumn(5, *args))

    def getVolume(self, *args):
        return int(self._getColumn(6, *args))

    def advanceDay(self): 
        # Jeżeli się da to odczytaj wskaźniki z kolejnego dnia i zwróć true.

        if self.currentDayIndex < self.lastDayIndex:
            self.currentDayIndex += 1
            self._currentDayRaw   = self._rawData[self.currentDayIndex]
            self.openPrice        = float(self._currentDayRaw[2])
            self.highPrice        = float(self._currentDayRaw[3])
            self.lowPrice         = float(self._currentDayRaw[4])
            self.closePrice       = float(self._currentDayRaw[5])
            self.volume           = int(self._currentDayRaw[6])
            self.OBV              = self.genOBV()
            return True

        return False

    def genOBV(self):
        if self.currentDayIndex == 0:
            return 0
        if self.closePrice > self.getClosePrice(-1):
            return self.OBV + self.volume
        if self.closePrice < self.getClosePrice(-1):
            return self.OBV - self.volume
        return self.OBV          

    # Pochodne wskaźniki

    def getRange(self, delta = 0):
        """
        Zakres zmian w danym dniu
        """
        delta = abs(delta)
        return self.getHighPrice(delta) - self.getLowPrice(delta);

    def getTR(self, delta = 0): 
        """
        Rzeczywisty zakres zmian w danym okresie to maksimum z następujących wartości:
        * maksymalna wartość - minimalna wartość
        * abs(poprzednie zamknięcie - maksymalna wartość)
        * abs(poprzednie zamknięcie - minimalna wartość)
        """
        delta = abs(delta)
        day = self.currentDayIndex-delta
        assert day >= 0
        return self._TR[day]

    def getATR(self, period, delta = 0):
        """
        n-dniowa EMA z wartości TR

        Typically, the Average True Range (ATR) is based on 14 periods and can
        be calculated on an intraday, daily, weekly or monthly basis. For this
        example, the ATR will be based on daily data. Because there must be a
        beginning, the first TR value is simply the High minus the Low, and the
        first 14-day ATR is the average of the daily TR values for the last 14
        days. After that, Wilder sought to smooth the data by incorporating the
        previous period's ATR value.

               
        Current ATR = [(Prior ATR x 13) + Current TR] / 14

          - Multiply the previous 14-day ATR by 13.
          - Add the most recent day's TR value.
          - Divide the total by 14
        """
        delta = abs(delta)
        day = self.currentDayIndex-delta
        assert day >= 0

        try:
            return self._ATR[period][day]
        except KeyError:
            self._generateATR(period)
            return self._ATR[period][day]

    def getPDM(self, delta = 0): 
        """
        +DM
        Positive Directional Movement
        """
        delta = abs(delta)
        day = self.currentDayIndex-delta
        assert day >= 0
        return self._PDM[day]

    def getNDM(self, delta = 0): 
        """
        -DM
        Negative Directional Movement
        """
        delta = abs(delta)
        day = self.currentDayIndex-delta
        assert day >= 0
        return self._NDM[day]

    def getPDI(self, period, delta = 0): 
        """
        +DI
        Positive Directional Indicator
        """

        delta = abs(delta)
        day = self.currentDayIndex-delta
        assert day >= 0
        try:
            return self._PDI[period][day]
        except KeyError:
            self._generatePDI(period)
            return self._PDI[period][day]

    def getNDI(self, period, delta = 0): 
        """
        -DI
        Negative Directional Indicator
        """

        delta = abs(delta)
        day = self.currentDayIndex-delta
        assert day >= 0
        try:
            return self._NDI[period][day]
        except KeyError:
            self._generateNDI(period)
            return self._NDI[period][day]

    def getADX(self, period, delta = 0): 
        """
        Average Directional Movement
        """

        delta = abs(delta)
        day = self.currentDayIndex-delta
        assert day >= 0
        try:
            return self._ADX[period][day]
        except KeyError:
            self._generateADX(period)
            return self._ADX[period][day]

    def getSMA(self, period, delta = 0): 
        """
        Simple Moving Average
        Średnia ostatnich `period' dni
        """

        """
        delta = abs(delta)
        sum = 0

        for d in xrange(delta,period+delta): 
            sum += self.getSharePrice(d)

        return sum / period
        """

        delta = abs(delta)
        day = self.currentDayIndex-delta
        assert day >= 0
        try:
            return self._SMA[period][day]
        except KeyError:
            self._generateSMA(period)
            return self._SMA[period][day]

    def getEMA(self, period, delta = 0): 
        """
        Exponential Moving Average -
        Średnia krocząca ważona na korzyść najnowszych dni 

        Multiplier = 2 / (periods+1)
        EMA = (close - EMA(previous)) * multiplier + EMA(previous)
        """

        delta = abs(delta)
        day = self.currentDayIndex-delta
        assert day >= 0
        try:
            return self._EMA[period][day]
        except KeyError:
            self._generateEMA(period)
            return self._EMA[period][day]

    def getMomentum(self, period, delta = 0):
        """
        Moment - 
        Różnica wartości zamknięcia dziś i `period' sesji temu, 
        """
        delta = abs(delta)
        return self.getClosePrice(delta) - self.getClosePrice(period+delta)

    def getROC(self, period, delta = 0): 
        """
        Rate of Change - procentowa zmiana ceny z dzisiejszej sesji, do ceny sprzed
        `period' sesji.
        """
        delta    = abs(delta)
        close    = self.getClosePrice(delta) 
        oldClose = self.getClosePrice(period+delta)

        return (close - oldClose) / float(oldClose)

    def getROV(self, period, delta = 0):
        """
        Rate of Volume - procentowa zmiana ceny z dzisiejszej sesji, do ceny sprzed
        `period' sesji.
        """
        delta    = abs(delta)
        volume    = self.getVolume(delta) 
        oldVolume = self.getVolume(period+delta)

        return (volume - oldVolume) / float(oldVolume)


    # {{{
    def __init__(self, symbol, lastSimulationDay = None, firstSimulationDay = 0):
        from csv import reader

        path = 'data/'+str.upper(symbol)+'.mst'
        reader = reader(open(path, 'r')) 

        # header = ["symbol", "date", "open", "high", "low", "close", "vol"]
        self._rawData = list(reader)[1:] # ucinamy nagłówek


        self.firstDayIndex    = 0
        if firstSimulationDay != None:
            self.firstDayIndex = firstSimulationDay

        self.currentDayIndex  = self.firstDayIndex-1
        self.lastDayIndex     = len(self._rawData)-1
        self.lastStockData    = len(self._rawData)-1

        if lastSimulationDay != None and lastSimulationDay != 0 and lastSimulationDay < self.lastDayIndex:
            self.lastDayIndex = lastSimulationDay

        assert (self.firstDayIndex < self.lastDayIndex)

        self._currentDayRaw   = self._rawData[self.currentDayIndex]
        self._TR              = []
        self._ATR             = {}
        self._SMA             = {}
        self._EMA             = {}
        self._PDM             = []
        self._NDM             = []
        self._PDI             = {}
        self._NDI             = {}
        self._ADX             = {}
        self._Highest         = {}
        self._Lowest          = {}
 
        self._generateDM()
        self._generateTR()
        self._generateOscilator()

    def _getHighest(self, delta, day):
        delta = abs(delta)
        high = 0
        assert day - delta >= 0
        for n in range(0, delta+1):
            tmp = float(self._rawData[day-n][4])
            if tmp > high:
               high = tmp
        return high

    def _getLowest(self, delta, day):
        delta = abs(delta)
        assert day-delta >= 0
        low = float(self._rawData[day][4])
        for n in range(1, delta+1):
            tmp = float(self._rawData[day-n][4])
            if tmp < low:
               low = tmp
        return low
            


    def _generateOscilator(self):
        for day in range(0, 7):
            low = []
            hight = []
            for n in range(0, 15+day):
                 low.append(-1)
                 hight.append(-1)
            highest = 0
            for l in range(15+day, self.lastDayIndex+1):
                 hight.append(self._getHighest(day, l))
                 low.append(self._getLowest(day, l))
            self._Highest[day] = hight
            self._Lowest[day] = low

    def _getRaw(self):
        return self._rawData

    def _getColumn(self, column, delta = 0): 

        if delta == 0:
            value = self._currentDayRaw[column]
        else:
            day   = self.currentDayIndex - abs(delta)
            #print "day: ", day
            assert (day >= 0 and day <= self.lastDayIndex), "incorrect day index"
            value = self._rawData[self.currentDayIndex - abs(delta)][column]

        return value

    def _generateTR(self): 

        tr = []

        data = self._rawData[0]
        high  = float(data[3])
        low   = float(data[4])

        tr.append(high-low) 

        old = data
        for fresh in self._rawData[1:]:
            high  = float(fresh[3])
            low   = float(fresh[4])
            close = float(old[5])

            tr.append(max(high, close) - min(low, close))

            old   = fresh

        self._TR = tr

    def _generateATR(self, period):

        # TODO: jest minimalne przesunięcie względem tego z Qtstalkera. zbadać czy
        #       to przez wybór początkowego, czy coś poważnego

        atr    = []
        alpha  = 1. / period

        self._generateTR()

        # pierwszy tr jest średnią początkowych tr, więc je sumujemy
        summ = 0
        for day in range(0, period):
            atr.append(-1)             # zaślepki dla łatwego indeksowania po dniach
            summ += self._TR[day]

        old = summ / period
        atr.append(old)

        for tr in self._TR[period:]:
            fresh   = tr * alpha + old * (1 - alpha)
            atr.append(fresh)
            old = fresh

        self._ATR[period] = atr

    def _generateSMA(self, period):

        from collections import deque

        sma    = []      # właściwe średnie
        values = deque() # ostatnie period wartości
        summ   = 0.0     # suma values 

        # zapełnij values
        for day in range(0, period):
            sma.append(-1)             # zaślepki dla łatwego indeksowania po dniach
            value = float(self._rawData[day][5])
            values.append(value)

        # mamy period wartości więc jest i pierwsza średnia 
        summ = sum(values)
        sma.append(summ / period)

        # dla wartości zamknięcia z kolejnych dni, dodawaj je do kolejki, wyrzucając stare,
        # uaktualniając przy tym średnią
        for data in self._rawData[period:]:
            x = float(data[5])
            values.append(x)
            summ  += x - values.popleft()
            sma.append(summ / period)

        self._SMA[period] = sma[1:]

    def _generateEMA(self, period):

        ema    = []
        alpha  = 2. / (period + 1)
        alpha_ = 1-alpha
    
        for day in range(0, period):
            ema.append(-1)

        self._generateSMA(period)
        new = self._SMA[period][period-1]
        ema.append(new)

        for data in self._rawData[period:]:
            close = float(data[5])
            new   = close * alpha + new * alpha_
            ema.append(new)

        self._EMA[period] = ema[1:]

    def _generateDM(self): 

        pdm = []
        mdm = []

        data    = self._rawData[0]
        oldhigh = float(data[3])
        oldlow  = float(data[4])

        pdm.append(0) 
        mdm.append(0) 

        for data in self._rawData[1:]:
            newhigh  = float(data[3])
            newlow   = float(data[4])

            upmove   = newhigh - oldhigh
            downmove = oldlow - newlow

            if upmove > downmove and upmove > 0:
                pdm.append(upmove)
            else:
                pdm.append(0)

            if downmove > upmove and downmove > 0:
                mdm.append(downmove)
            else:
                mdm.append(0)

            oldlow, oldhigh = newlow, newhigh

        self._PDM = pdm
        self._NDM = mdm

    def _generatePDI(self, period):

        # oblicz EMA z +DM
        epdm   = []
        alpha  = 1. / period

        self._generateDM()

        summ = 0
        for day in range(0, period):
            epdm.append(-1)             # zaślepki dla łatwego indeksowania po dniach
            summ += self._PDM[day]

        old = summ / period
        epdm.append(old)

        for pdm in self._PDM[period:]:
            fresh   = pdm * alpha + old * (1 - alpha)
            epdm.append(fresh)
            old = fresh

        self._generateATR(period)

        pdi = []
        pdi.append(-1)
        for epdm_, atr_ in zip(epdm[1:], self._ATR[period][1:]): 
            if epdm == -1 or atr_ == -1:
                pdi.append(-1)
            else:
                pdi.append(100*epdm_/ atr_)

        self._PDI[period] = pdi

    def _generateNDI(self, period):

        # oblicz EMA z -DM
        endm   = []
        alpha  = 1. / period

        self._generateDM()

        summ = 0
        for day in range(0, period):
            endm.append(-1)             # zaślepki dla łatwego indeksowania po dniach
            summ += self._NDM[day]

        old = summ / period
        endm.append(old)

        for ndm in self._NDM[period:]:
            fresh   = ndm * alpha + old * (1 - alpha)
            endm.append(fresh)
            old = fresh

        self._generateATR(period)

        pdi = []
        pdi.append(-1)
        for endm_, atr_ in zip(endm[1:], self._ATR[period][1:]): 
            if endm == -1 or atr_ == -1:
                pdi.append(-1)
            else:
                pdi.append(100*endm_/ atr_)

        self._NDI[period] = pdi

    def _generateADX(self, period):

        adx    = []
        alpha  = 1. / period

        self._generatePDI(period)
        self._generateNDI(period)

        for day in range(0, period):
            adx.append(-1)
        summ = 0

        for day in range(period, 2*period):
            adx.append(-1)
            pdi_ = self._PDI[period][day]
            ndi_ = self._NDI[period][day]

            summ += (100. * abs(pdi_ - ndi_)) / (pdi_ + ndi_)


        old = summ / period
        adx.append(old)

        for pdi_, ndi_ in zip(self._PDI[period][2*period:], self._NDI[period][2*period:]):
            new =  (100. * abs(pdi_ - ndi_)) / (pdi_ + ndi_)

            fresh   = new * alpha + old * (1 - alpha)
            adx.append(fresh)
            old = fresh

        self._ADX[period] = adx
    # }}}
# }}} 

if __name__ == "__main__":

    egb = StockModel('cerabud')

    i = 0
    period = 14
    delta = 0

    while (egb.advanceDay()):
        i += 1
        if (egb.firstDayIndex + period + delta  <= egb.currentDayIndex):
            """
            print "%3d| %s | O%.2f H%.2f L%.2f C%.2f @ SMA %.5f" % (i,
                                                           egb._getColumn(1),
                                                           egb.getOpenPrice(),
                                                           egb.getHighPrice(),
                                                           egb.getLowPrice(),
                                                           egb.getClosePrice(),
                                                           egb.getSMA(period)
                                                            )
            """
            print("%d %.5f" % (i, egb.getMomentum(period)))

# vim: fdm=marker
